from torchvision import transforms as T
from pathlib import Path
import torch

transformer = T.Compose([
    T.Resize((256, 256), antialias=True),
    T.CenterCrop((224, 224)),
    T.ToTensor(),   # 将图像转换为0-1范围
    # 图像标准化，依据ImageNet数据集的均值和标准差
    T.Normalize(mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225])
])

"""
当我们去加载自己准备的图像数据集的时候，可以使用
1、os + DataSet方法，制作元数据+自定义数据集
2、高级API ImageFolder
"""
from torchvision.datasets import ImageFolder

# 数据目录解析：优先使用项目根目录的 natural_images，其次回退同级目录
base_dir = Path(__file__).resolve().parent
candidates = [
    base_dir.parent / "natural_images",  # 项目根目录外部共享数据集
    base_dir / "natural_images",         # 回退：同级目录
]
data_dir = next((p for p in candidates if p.exists()), candidates[0])

# 使用 ImageFolder 读取整个数据集后进行随机划分（例如 8:2）
full_ds = ImageFolder(
    root=str(data_dir),
    transform=transformer
)
num_total = len(full_ds)
val_ratio = 0.2
num_val = int(num_total * val_ratio)
num_train = num_total - num_val
generator = torch.Generator().manual_seed(42)  # 固定随机种子，保证可复现
train_ds, val_ds = torch.utils.data.random_split(full_ds, [num_train, num_val], generator=generator)

""" 数据集加载器 """
from torch.utils.data import DataLoader, random_split
train_dl = DataLoader(train_ds, batch_size=64, shuffle=True)
val_dl = DataLoader(val_ds, batch_size=128, shuffle=False)

""" 标签映射表 """
# 使用完整数据集的类别映射（与 ImageFolder 的字母序一致）
idx2label = {idx: label for label, idx in full_ds.class_to_idx.items()}

if __name__ == "__main__":
    print("训练集的大小", len(train_ds))
    img, label = train_ds[1000]
    print(img.shape, label)
    print("列表标签", train_ds.classes)
    print("总类别数", len(train_ds.classes))
    label2idx = train_ds.class_to_idx
    print("标签 to 索引", label2idx)
    idx2label = {idx:label for label,idx in label2idx.items()}
    print("标签映射表", idx2label)

    for imgs, labels in train_dl:
        print(imgs.shape)
        print(labels)
        break